from omegaconf import DictConfig

import torch
from torchvision.transforms.v2 import RandomRotation, RandomHorizontalFlip, RandomResizedCrop, RandomApply, Compose

from ecgcmr.imaging.img_augmentations.BrightnessContrastSaturation import CustomColorJitter
from ecgcmr.imaging.img_augmentations.GaussianBlur import CustomGaussianBlur
from ecgcmr.imaging.img_augmentations.SampleTime import SampleTime, SampleMRI
from ecgcmr.imaging.img_augmentations.GaussianNoise import CustomGaussianNoise

from ecgcmr.multimodal.multimodal_utils.ToTensor import ToTensor
from ecgcmr.multimodal.multimodal_utils.SqueezeUnsqueeze import UnsqueezeDimension


class ImageAugmentations:
    def __init__(self,
                 cfg: DictConfig,
                 img_size: int = 84,
                 apply_augmentations: bool = True,
                 use_peaks_location: bool = False
                 ) -> None:
        self.cfg = cfg
        self.img_size = img_size
        self.apply_augmentations = apply_augmentations
        self.use_peaks_location = use_peaks_location

        self.img_augmentations_steps = self._create_augmentations()

    def _check_cfg_enabled(self, config_path: str) -> bool:
        """
        Utility to check if a config option is enabled.
        Returns `False` if the config does not exist or is not explicitly set to `True`.
        """
        return self.cfg.get(config_path, {}).get('enable', False)

    def _create_augmentations(self) -> Compose:
        augmentations = []

        self._add_time_sample(augmentations)
        self._add_to_tensor(augmentations)

        if self.apply_augmentations:
            self._add_gaussian_noise(augmentations)
            self._add_gaussian_blur(augmentations)
            self._add_color_jitter(augmentations)
            self._add_random_resized_crop(augmentations)
            self._add_random_horizontal_flip(augmentations)
            self._add_random_rotation(augmentations)

        self._add_unsqueeze_dimension(augmentations)

        return Compose(augmentations)

    def _add_time_sample(self, augmentations: list) -> None:
        if self._check_cfg_enabled('time_sample'):
            augmentations.append(SampleMRI(result_n_frames=self.cfg.time_sample.result_n_frames))

    def _add_to_tensor(self, augmentations: list) -> None:
        augmentations.append(ToTensor())

    def _add_random_resized_crop(self, augmentations: list) -> None:
        if self._check_cfg_enabled('random_resized_crop'):
            augmentations.append(
                RandomResizedCrop(
                    size=self.img_size,
                    scale=(self.cfg.random_resized_crop.scale_from, self.cfg.random_resized_crop.scale_to),
                    ratio=(self.cfg.random_resized_crop.ratio_from, self.cfg.random_resized_crop.ratio_to)
                )
            )

    def _add_random_horizontal_flip(self, augmentations: list) -> None:
        if self._check_cfg_enabled('random_horizontal_flip'):
            augmentations.append(
                RandomHorizontalFlip(p=self.cfg.random_horizontal_flip.prob)
            )

    def _add_random_rotation(self, augmentations: list) -> None:
        if self._check_cfg_enabled('random_rotation'):
            augmentations.append(
                RandomApply(
                    transforms=[RandomRotation(degrees=self.cfg.random_rotation.degrees)],
                    p=self.cfg.random_rotation.prob
                )
            )

    def _add_color_jitter(self, augmentations: list) -> None:
        if self._check_cfg_enabled('color_jitter'):
            augmentations.append(
                RandomApply(
                    transforms=[
                        CustomColorJitter(
                            strength=self.cfg.color_jitter.strength,
                            brightness=self.cfg.color_jitter.brightness,
                            contrast=self.cfg.color_jitter.contrast,
                        )
                    ],
                    p=self.cfg.color_jitter.prob
                )
            )

    def _add_gaussian_blur(self, augmentations: list) -> None:
        if self._check_cfg_enabled('gaussian_blur'):
            augmentations.append(
                RandomApply(
                    transforms=[CustomGaussianBlur(img_size=self.img_size,
                                                   min_strength=self.cfg.gaussian_blur.min_strength,
                                                   max_strength=self.cfg.gaussian_blur.max_strength)],
                    p=self.cfg.gaussian_blur.prob
                )
            )
    def _add_gaussian_noise(self, augmentations: list) -> None:
        if self._check_cfg_enabled('gaussian_noise'):
            augmentations.append(
                RandomApply(
                    transforms=[CustomGaussianNoise(min_strength=self.cfg.gaussian_noise.min_strength,
                                                    max_strength=self.cfg.gaussian_noise.max_strength)],
                    p=self.cfg.gaussian_noise.prob
                )
            )

    def _add_unsqueeze_dimension(self, augmentations: list) -> None:
        if self._check_cfg_enabled('unsqueeze'):
            augmentations.append(UnsqueezeDimension(position=self.cfg.unsqueeze.position))

    def __call__(self, image_data: dict) -> torch.Tensor:
        return self.img_augmentations_steps(image_data)
